/**
 * @file uring_server.cc
 * @author SHADOW (shadow_yuan@qq.com)
 * @brief tcp server demo (io_uring)
 * @version 0.1
 * @date 2024-10-11
 *
 */
#ifdef _WIN32
#include <WinSock2.h>
#include <WS2tcpip.h>
#include <MSWSock.h>
#else
#include <arpa/inet.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <liburing.h>
#endif // _WIN32
#include <string.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <inttypes.h>
#include <errno.h>

#include <atomic>
#include <chrono>
#include <string>
#include <iostream>
#include <memory>

#ifndef _WIN32
typedef int SOCKET;
#define CLOSESOCKET close
#else
#define CLOSESOCKET closesocket
#pragma comment(lib, "ws2_32.lib")
#endif // !_WIN32

typedef std::atomic<intptr_t> SpinLock;

constexpr int kNullEvent = 0;
constexpr int kAcceptEvent = 1;
constexpr int kConnectEvent = 1 << 1;
constexpr int kSendEvent = 1 << 2;
constexpr int kRecvEvent = 1 << 3;

struct DataNode {
    int event;
    int size;
    SOCKET fd;
    char data[];
};

struct DataNodeAccept {
    int event;
    int size;
    SOCKET fd;
    struct sockaddr_in addr;
    socklen_t addr_len;
};

constexpr int kEmptyNodeSize = sizeof(DataNode);
constexpr int kMaxBuffer = 1024;

#define DataNodePtr(node) (&node->data[0])

static inline DataNodeAccept *MakeDataNodeAccept(SOCKET fd) {
    char *buff = new char[sizeof(DataNodeAccept)];
    auto *node = reinterpret_cast<DataNodeAccept *>(buff);
    node->event = kAcceptEvent;
    node->size = sizeof(DataNodeAccept) - kEmptyNodeSize;
    node->addr_len = sizeof(node->addr);
    node->fd = fd;
    return node;
}

static inline DataNode *MakeDataNode(int size, int event, SOCKET fd) {
    char *buff = new char[kEmptyNodeSize + size];
    DataNode *node = reinterpret_cast<DataNode *>(buff);
    node->event = event;
    node->size = size;
    node->fd = fd;
    return node;
}

static inline void FreeDataNode(DataNode *node) { delete[] (char *)node; }

int set_socket_nonblocking(SOCKET fd, int nonblocking) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags < 0) return -1;
    if (nonblocking) {
        flags |= O_NONBLOCK;
    } else {
        flags &= ~O_NONBLOCK;
    }
    return fcntl(fd, F_SETFL, flags);
}

struct io_uring g_ring;
bool g_running = true;

void AsyncAccept(SOCKET fd) {
    auto node = MakeDataNodeAccept(fd);
    auto sqe = io_uring_get_sqe(&g_ring);
    io_uring_prep_accept(sqe, fd, (struct sockaddr *)&node->addr, &node->addr_len,
                         SOCK_NONBLOCK | SOCK_CLOEXEC);
    io_uring_sqe_set_data(sqe, node);
    io_uring_submit(&g_ring);
    printf("submit async accept, listen_fd: %d\n", fd);
}

void AsyncRecv(SOCKET fd) {
    auto node = MakeDataNode(kMaxBuffer, kRecvEvent, fd);
    struct io_uring_sqe *sqe = io_uring_get_sqe(&g_ring);
    io_uring_prep_recv(sqe, fd, DataNodePtr(node), node->size, 0);
    io_uring_sqe_set_data(sqe, node);
    io_uring_submit(&g_ring);
}

void AsyncSend(SOCKET fd, const std::string &text) {
    auto node = MakeDataNode(text.size(), kSendEvent, fd);
    memcpy(DataNodePtr(node), text.c_str(), text.size());
    struct io_uring_sqe *sqe = io_uring_get_sqe(&g_ring);
    io_uring_prep_send(sqe, fd, DataNodePtr(node), node->size, 0);
    io_uring_sqe_set_data(sqe, node);
    io_uring_submit(&g_ring);
}

std::string get_client_address(DataNodeAccept *node) {
    char ip[64] = {0};
    int port = ntohs(node->addr.sin_port);
    inet_ntop(AF_INET, &node->addr.sin_addr, ip, sizeof(ip));

    char buff[128] = {0};
    int len = snprintf(buff, sizeof(buff) - 1, "%s:%d", ip, port);
    return std::string(buff, len);
}

void MainEventLoop() {
    while (g_running) {
        io_uring_cqe *cqe = NULL;
        int ret = io_uring_wait_cqe(&g_ring, &cqe);
        if (ret != 0) {
            printf("io_uring_wait_cqe return %d, cqe: %p\n", ret, cqe);
            if (cqe) {
                DataNode *node = (DataNode *)io_uring_cqe_get_data(cqe);
                io_uring_cqe_seen(&g_ring, cqe);
                FreeDataNode(node);
            }
            continue;
        }

        DataNode *node = (DataNode *)io_uring_cqe_get_data(cqe);
        switch (node->event) {
        case kAcceptEvent: {
            auto obj = (DataNodeAccept *)node;
            if (cqe->res != 0) {
                std::string peer = get_client_address(obj);
                printf("accept event, [success] listen_fd: %d, client(%d) addr:[%s]\n", obj->fd,
                       cqe->res, peer.c_str());
                set_socket_nonblocking(cqe->res, 1);
                AsyncRecv(cqe->res);
            } else {
                printf("accept event, [failure] listen_fd: %d\n", obj->fd);
            }
            AsyncAccept(obj->fd);
            break;
        }
        case kSendEvent:
            printf("send event, written %d bytes\n", cqe->res);
            break;
        case kRecvEvent:
            if (cqe->res == 0) {
                printf("recv event, client(%d) closed!\n", node->fd);
                CLOSESOCKET(node->fd);
            } else {
                std::string text(DataNodePtr(node), cqe->res);
                printf("recv event, recv %d bytes '%s' from client(%d).\n", cqe->res, text.c_str(),
                       node->fd);
                AsyncSend(node->fd, text);
                AsyncRecv(node->fd);
            }
            break;
        }
        io_uring_cqe_seen(&g_ring, cqe);
        // io_uring_peek_cqe(&g_ring, &cqe);
        FreeDataNode(node);
    }
}

int main(int argc, char *argv[]) {
    struct sockaddr_in server_addr;
    SOCKET listen_fd;
    int port;

    if (argc != 2) {
        printf("usage: server <port>\n");
        return -1;
    }

#ifdef _WIN32
    WSADATA data;
    int status;
    status = WSAStartup(MAKEWORD(2, 0), &data);
    if (status != 0) {
        printf("failed to init win sock\n");
        return -1;
    }
#endif // _WIN32

    port = atoi(argv[1]);
    printf("using localhost:%d\n", port);

    // init io_uring object
    io_uring_queue_init(64, &g_ring, 0);

    if ((listen_fd = socket(AF_INET, SOCK_STREAM, 0)) == (SOCKET)(-1)) {
        printf("failed to create socket\n");
        return -1;
    }

    set_socket_nonblocking(listen_fd, 1);

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    server_addr.sin_addr.s_addr = htonl(INADDR_ANY);

    bind(listen_fd, (struct sockaddr *)&server_addr, sizeof(server_addr));

    listen(listen_fd, 128);

    AsyncAccept(listen_fd);

    MainEventLoop();

    CLOSESOCKET(listen_fd);
#ifdef _WIN32
    WSACleanup();
#endif // _WIN32
    return 0;
}
